"""
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    created by lane.chang on '28/05/2024'
    comment: 知识库相关
"""
import time
from fastapi import APIRouter, Request, Query, Body
from fastapi.responses import StreamingResponse
from typing import Optional

from config import Config
from project import json_response
from project.lib.common import get_redis
from project.lib.decorator import api_common, exec_duration
from project.service.knowledge import KnowledgeService
from project.model.basic import KlItem
from project.model.basic import ChatMessage


knowledge_router = APIRouter(prefix='', tags=['知识库'])


@knowledge_router.get('/knowledge/info', summary='知识库信息')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def get_knowledge_info(request: Request, key: str = Query(..., title='知识库key'), session=None):
    """
    **desc**: 取得指定知识库信息

    **url**: http(s)://{0}.udicaria.com/knowledge/info

    **method**: GET

    **params**:

        <query>:
            key: 知识库key

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data:
            exec_time: 接口后端处理时长
        }
    """
    file_object = KnowledgeService.get_knowledge_file_object(session=session, knowledge_key=key)

    return StreamingResponse(file_object, media_type="text/plain", headers={"Content-Disposition": f"attachment; filename=knowledge_{key}.json"})


@knowledge_router.post('/knowledge/retrieve', summary='知识库匹配')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def retrieve_knowledge(request: Request,
                             key: str = Body(..., title='知识库key'),
                             user_message: str = Body(..., title='用户话术'),
                             count: int = Body(3, title='取出条数'),
                             session=None):
    """
    **desc**: 取得知识库匹配信息

    **url**: http(s)://{0}.udicaria.com/knowledge/retrieve

    **method**: GET

    **params**:

        <body>:
            key: 知识库key
            user_message: 用户话术
            count: 匹配条数 默认3条

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data:
            exec_time: 接口后端处理时长
        }
    """
    data = KnowledgeService.retrieve(session=session, organization_code=key, text=user_message, count=count)

    return json_response(data=data)


@knowledge_router.post('/knowledge/flush/streaming', summary='知识库刷新')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def flush_knowledge_streaming(request: Request,
                          key: str = Body(..., title='知识库key'),
                          kl_items: list[KlItem] = Body(..., title='知识库'),
                          session=None):
    """
    **desc**: 知识库刷新(流式)

    **url**: http(s)://{0}.udicaria.com/knowledge/flush

    **method**: POST

    **params**:

        <body>:
            key: 知识库key
            user_message: 用户话术
            kl_items: 知识库项目列表{'question': xxx, 'answer': xxx}

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data:
            exec_time: 接口后端处理时长
        }
    """
    return StreamingResponse(KnowledgeService.flush_knowledge_streaming(session=session, knowledge_key=key, kls=kl_items),
                             media_type="application/octet-stream",
                             headers={"Content-Disposition": f"attachment; filename=flush_result_{int(time.time())}.text"})


@knowledge_router.get('/knowledge/flush/rwt/streaming', summary='知识库刷新')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def flush_rwt_knowledge_streaming(request: Request, session=None):
    """
    **desc**: 如味堂知识库刷新(流式)

    **url**: http(s)://{0}.udicaria.com/knowledge/flush/rwt/streaming

    **method**: GET

    **params**:

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data:
            exec_time: 接口后端处理时长
        }
    """
    key = 'knowledge_flush_rwt_streaming'
    rds = get_redis()
    if rds.exists(key):
        return json_response(data='知识库更新中, 请耐心等待')
    else:
        return StreamingResponse(KnowledgeService.flush_rwt_knowledge_streaming(session=session),
                                 media_type="application/octet-stream",
                                 headers={"Content-Disposition": f"attachment; filename=flush_result_{int(time.time())}.text"})


@knowledge_router.post('/knowledge/flush', summary='知识库刷新')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def flush_knowledge(request: Request,
                          knowledge_detail_code: Optional[str] = Body(None, title='知识库详情编号'),
                          organization_code: str = Body(..., title='机构编号'),
                          text: str = Body(..., title='知识库文本'),
                          session=None):
    """
    **desc**: 知识库刷新

    **url**: http(s)://{0}.udicaria.com/knowledge/flush

    **method**: POST

    **params**:

        <body>:
            knowledge_detail_code: 知识库详情编号
            organization_code: 机构编号
            agent_code: 智能体编号
            text: 知识库文本

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data:
            exec_time: 接口后端处理时长
        }
    """
    await KnowledgeService.flush_knowledge(session, organization_code, text, knowledge_detail_code)

    return json_response()


@knowledge_router.post('/knowledge/delete', summary='知识库删除')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def delete_knowledge(request: Request,
                           knowledge_detail_code: str = Body(..., title='知识库详情编号', embed=True),
                           session=None):
    """
    **desc**: 知识库删除

    **url**: http(s)://{0}.udicaria.com/knowledge/delete

    **method**: POST

    **params**:

        <body>:
            knowledge_detail_code: 知识库详情编号

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data:
            exec_time: 接口后端处理时长
        }
    """
    await KnowledgeService.delete_knowledge(session, knowledge_detail_code)

    return json_response()


@knowledge_router.post('/knowledge/delete/batch', summary='知识库删除(批量)')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def delete_knowledge_batch(request: Request,
                                 knowledge_detail_codes: list[str] = Body(..., title='知识库详情编号列表', embed=True),
                                 session=None):
    """
    **desc**: 知识库删除

    **url**: http(s)://{0}.udicaria.com/knowledge/delete/batch

    **method**: POST

    **params**:

        <body>:
            knowledge_detail_codes: 知识库详情编号

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data:
            exec_time: 接口后端处理时长
        }
    """
    await KnowledgeService.delete_knowledge_batch(session, knowledge_detail_codes)

    return json_response()


@knowledge_router.post('/knowledge/match/list', summary='')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def get_knowledge_match_list(request: Request,
                                   organization_code: Optional[str] = Body(None, title='机构编号'),
                                   knowledge_detail_codes: Optional[list[str]] = Body(None, title='知识库详情编号列表'),
                                   text: str = Body(..., title='文本'),
                                   count: int = Body(3, title='返回条数'),
                                   session=None):
    """
    **desc**: 知识库匹配列表

    **url**: http(s)://{0}.udicaria.com/knowledge/match/list

    **method**: POST

    **params**:

        <body>:
            organization_code: 机构编号
            knowledge_detail_codes: 知识库详情编号列表
            text: 文本
            count: 返回条数 默认3

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data: [
                {
                    knowledge_detail_code: 知识库编号
                    text: 文本
                }
            ]
            exec_time: 接口后端处理时长
        }
    """
    data = await KnowledgeService.get_knowledge_match_list(session, text, organization_code, knowledge_detail_codes, count)

    return json_response(data=data)


@knowledge_router.post('/knowledge/chat/streaming', summary='')
@api_common(doc_replace=(('{0}', Config.HTTP_DOC_FIX),))
@exec_duration(add_return=True)
async def do_knowledge_chat_streaming(request: Request,
                                      organization_code: str = Body(None, title='机构编号'),
                                      knowledge_detail_codes: list[str] = Body(None, title='知识库详情编号列表'),
                                      user_message: str = Body(..., title='用户话术'),
                                      chat_messages: list[ChatMessage] = Body([], title='对话历史'),
                                      session=None):
    """
    **desc**: 知识库聊天(流式)

    **url**: http(s)://{0}.udicaria.com/knowledge/chat/streaming

    **method**: POST

    **params**:

        <body>:
            organization_code: 机构编号
            knowledge_detail_codes: 知识库详情编号列表
            user_message: 文本
            chat_messages: 对话历史

    **return**:json数据结构::

        {
            code: 0(成功)/非0(失败)
            message: 成功返回success, 否则返回error信息
            data: [
                {
                    knowledge_detail_code: 知识库编号
                    text: 文本
                }
            ]
            exec_time: 接口后端处理时长
        }
    """
    response_stream = await KnowledgeService.do_knowledge_chat_streaming(session, user_message, organization_code, knowledge_detail_codes, chat_messages=chat_messages)

    return StreamingResponse(response_stream, media_type='text/plain')
